import pandas as pd
import os
import sys
import traceback
import logging
import concurrent.futures
import subprocess
from PyQt6 import QtWidgets, QtCore
from openpyxl import load_workbook
import chardet
from PyQt6.QtWidgets import QTableWidget, QTableWidgetItem, QPushButton, QHeaderView, QListWidget, QVBoxLayout, QHBoxLayout

class IPCheckerApp(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.checking = False
        self.selected_files = []
        self.export_excel = True  # 默认导出Excel

    def init_ui(self):
        self.setWindowTitle('IP地址检测器')
        self.setGeometry(100, 100, 1000, 800)  # 加大窗口尺寸

        self.layout = QtWidgets.QVBoxLayout()

        # 文件选择部分
        file_section = QtWidgets.QVBoxLayout()
        
        # 按钮和选项布局
        top_layout = QtWidgets.QHBoxLayout()
        
        # 文件按钮布局
        file_button_layout = QtWidgets.QVBoxLayout()
        self.file_button = QtWidgets.QPushButton('选择文件')
        self.clear_files_button = QtWidgets.QPushButton('清除文件列表')
        file_button_layout.addWidget(self.file_button)
        file_button_layout.addWidget(self.clear_files_button)
        top_layout.addLayout(file_button_layout)
        
        # 导出选项
        export_layout = QtWidgets.QVBoxLayout()
        self.export_checkbox = QtWidgets.QCheckBox('导出到Excel')
        self.export_checkbox.setChecked(True)
        self.export_checkbox.toggled.connect(self.toggle_export)
        export_layout.addWidget(self.export_checkbox)
        top_layout.addLayout(export_layout)
        
        file_section.addLayout(top_layout)

        # 文件列表显示
        self.file_list = QListWidget()
        self.file_list.setMaximumHeight(80)  # 限制文件列表高度
        file_section.addWidget(self.file_list)
        
        self.layout.addLayout(file_section)

        # 进度条
        self.progress = QtWidgets.QProgressBar()
        self.progress.setVisible(False)
        self.layout.addWidget(self.progress)

        # 按钮布局
        button_layout = QtWidgets.QHBoxLayout()
        self.start_button = QtWidgets.QPushButton('开始检测')
        self.start_button.setEnabled(False)
        self.clear_button = QtWidgets.QPushButton('清除结果')
        self.clear_button.setEnabled(False)
        button_layout.addWidget(self.start_button)
        button_layout.addWidget(self.clear_button)
        self.layout.addLayout(button_layout)

        # 结果表格
        self.result_table = QTableWidget()
        self.result_table.setColumnCount(4)  # 增加到4列
        self.result_table.setHorizontalHeaderLabels(['IP地址', '状态', '描述', '详细信息'])
        header = self.result_table.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
        header.setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)
        header.setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch)
        header.setSectionResizeMode(3, QHeaderView.ResizeMode.Stretch)
        self.layout.addWidget(self.result_table)

        # 状态栏
        self.status_label = QtWidgets.QLabel('')
        self.layout.addWidget(self.status_label)

        self.setLayout(self.layout)

        # 连接信号
        self.file_button.clicked.connect(self.select_file)
        self.clear_files_button.clicked.connect(self.clear_file_list)
        self.start_button.clicked.connect(self.start_checking)
        self.clear_button.clicked.connect(self.clear_results)

    def toggle_export(self, state):
        self.export_excel = state

    def select_file(self):
        try:
            file_dialog = QtWidgets.QFileDialog()
            files, _ = file_dialog.getOpenFileNames(
                self,
                "选择包含IP地址的Excel文件",
                "",
                "Excel files (*.xlsx *.xls)"
            )
            if files:
                for file_path in files:
                    if file_path not in self.selected_files:
                        self.selected_files.append(file_path)
                        self.file_list.addItem(os.path.basename(file_path))
                self.start_button.setEnabled(True)
                self.update_status()
        except Exception as e:
            self.show_error("选择文件时发生错误", str(e))

    def clear_file_list(self):
        self.selected_files.clear()
        self.file_list.clear()
        self.start_button.setEnabled(False)
        self.update_status()

    def update_status(self):
        count = len(self.selected_files)
        self.status_label.setText(f"已选择 {count} 个文件")

    def start_checking(self):
        if self.checking:
            return

        try:
            self.checking = True
            self.disable_controls()
            self.backup_and_check()
        except Exception as e:
            self.show_error("检测过程发生错误", str(e))
        finally:
            self.checking = False
            self.enable_controls()

    def backup_and_check(self):
        if not self.selected_files:
            self.show_error("文件错误", "请选择有效的Excel文件")
            return

        try:
            total_ips = []
            self.df_dict = {}  # 存储所有文件的DataFrame

            # 首先读取所有文件的IP
            for file_path in self.selected_files:
                if not os.path.exists(file_path):
                    self.show_error("文件错误", f"文件不存在: {file_path}")
                    continue

                # 备份文件
                backup_file_path = file_path + '.bak'
                os.replace(file_path, backup_file_path)

                # 读取IP列表
                ip_list, df, descriptions = read_ips_from_excel(backup_file_path)
                if ip_list:
                    total_ips.extend(zip(ip_list, descriptions))
                    self.df_dict[file_path] = df

            if not total_ips:
                self.show_error("数据错误", "所有文件中都没有找到IP地址")
                return

            # 设置进度条
            self.results = {}
            self.progress.setVisible(True)
            self.progress.setMaximum(len(total_ips))
            self.progress.setValue(0)

            self.result_table.setRowCount(0)
            self.status_label.setText("正在检测中...")

            # 开始检测
            max_workers = min(50, os.cpu_count() * 5)
            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_to_ip = {executor.submit(ping_ip, ip): (ip, desc) for ip, desc in total_ips}
                completed = 0
                for future in concurrent.futures.as_completed(future_to_ip):
                    ip, desc = future_to_ip[future]
                    completed += 1
                    self.progress.setValue(completed)
                    try:
                        status = future.result()
                        status_text = "已连接" if status else "未连接"
                        self.results[ip] = status_text
                        self.add_result_row(ip, status_text, desc)
                    except Exception as exc:
                        self.add_result_row(ip, "错误", desc, str(exc))

            # 根据选项决定是否更新Excel文件
            if self.export_excel:
                for file_path, df in self.df_dict.items():
                    self.update_excel_file(file_path, df)
                self.status_label.setText("检测完成，结果已保存到Excel文件")
            else:
                self.status_label.setText("检测完成")

        except Exception as e:
            self.show_error("处理过程出错", str(e))
        finally:
            self.progress.setVisible(False)

    def show_error(self, title, message):
        QtWidgets.QMessageBox.critical(self, title, message)

    def disable_controls(self):
        self.file_button.setEnabled(False)
        self.start_button.setEnabled(False)
        self.clear_button.setEnabled(False)

    def enable_controls(self):
        self.file_button.setEnabled(True)
        self.start_button.setEnabled(True)
        self.clear_button.setEnabled(bool(self.result_table.rowCount()))

    def update_excel_file(self, file_path, df):
        try:
            for ip, status in self.results.items():
                if 'IP地址' in df.columns and ip in df['IP地址'].values:
                    df.loc[df['IP地址'] == ip, '网络状态'] = status

            writer = pd.ExcelWriter(file_path, engine='openpyxl')
            df.to_excel(writer, index=False)
            worksheet = writer.sheets['Sheet1']
            for col, width in zip(['A', 'B', 'C', 'D', 'E'], [10, 15, 30, 15, 30]):
                worksheet.column_dimensions[col].width = width
            writer.close()
        except Exception as e:
            self.show_error(f"保存文件错误 ({file_path})", str(e))

    def clear_results(self):
        self.result_table.setRowCount(0)
        self.status_label.setText('')
        self.clear_button.setEnabled(False)

    def add_result_row(self, ip, status, description='', details=''):
        row = self.result_table.rowCount()
        self.result_table.insertRow(row)
        self.result_table.setItem(row, 0, QTableWidgetItem(ip))
        self.result_table.setItem(row, 1, QTableWidgetItem(status))
        self.result_table.setItem(row, 2, QTableWidgetItem(description))
        self.result_table.setItem(row, 3, QTableWidgetItem(details))

def read_ips_from_excel(file_path):
    # 读取Excel文件中的IP地址
    df = pd.read_excel(file_path, engine='openpyxl')
    ip_column = next((col for col in ['IP地址', 'IP'] if col in df.columns), None)
    if not ip_column:
        raise KeyError("Excel文件中未找到'IP地址'或'IP'列，请检查列名。")
    
    descriptions = []
    if '描述' in df.columns:
        descriptions = df['描述'].fillna('').tolist()
    else:
        descriptions = [''] * len(df[ip_column])
        
    return df[ip_column].tolist(), df, descriptions

def ping_ip(ip):
    # 使用ping命令检测IP地址的网络通断情况
    try:
        result = subprocess.run(['ping', '-n', '1', ip], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        output = result.stdout
        detected_encoding = chardet.detect(output)['encoding']
        output = output.decode(detected_encoding or 'utf-8', errors='ignore')
        return "TTL=" in output
    except subprocess.CalledProcessError:
        return False
    except Exception as exc:
        print(f"{ip} generated an exception: {exc}")
        return False

# 设置日志记录
def setup_logging():
    log_file = os.path.join(os.path.dirname(__file__), 'ip_checker.log')
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler(sys.stdout)
        ]
    )
    return logging.getLogger(__name__)

if __name__ == "__main__":
    try:
        logger = setup_logging()
        logger.info("启动应用程序")
        
        app = QtWidgets.QApplication([])
        window = IPCheckerApp()
        window.show()
        
        # 确保窗口显示在最前面
        window.raise_()
        window.activateWindow()
        
        sys.exit(app.exec())
    except Exception as e:
        error_msg = f"程序发生错误:\n{str(e)}\n{traceback.format_exc()}"
        logging.error(error_msg)
        
        # 如果GUI还没有初始化，使用命令行显示错误
        print(error_msg)
        
        # 保持窗口打开
        input("按回车键退出...")
        sys.exit(1)